library(LDlinkR)
library(SeqArray)
library(SNPRelate)
library(ggpubr)
library(tidyverse)

api_token <- Sys.getenv("LDLINK_TOKEN")

all.height.snp <- read_csv('mission/exercise/ensembl.height.snp.csv') |>
  pull(id)

ighg1.ld <- c(all.height.snp, 'rs117518546') |>
  LDmatrix(pop = 'EAS',
           token = api_token,
           genome_build = 'grch38_high_coverage')

ighg1.ld |>
  as_tibble() |>
  select(RS_number, rs117518546) |>
  slice_max(rs117518546, n=2)

# calc Fst between 2 pop ---------
#' Calculate Fst value between 2 population
#'
#' @param p1 float, allele freq of population 1
#' @param p2 float, allele freq of population 2
#' @param c1 integer, total number of population 1
#' @param c2 integer, total number of population 2
#'
#' @return Return a float between 0-1. 0 means no differentiation between population 1 & 2. 1 means total differentiation.
#' @export 
#'
#' @examples
calc_fst2 <- function(p1, p2, c1, c2){
  calc_hetero <- function(p) 2*p*(1-p)
  p0 <- (c1*p1 + c2*p2)/(c1+c2)
  divrs0 <- calc_hetero(p0)
  divrs1 <- calc_hetero(p1)
  divrs2 <- calc_hetero(p2)
  1 - (divrs1 + divrs2) / 2 / divrs0
}

# 1KGP sample-subpop-pop meta -----------
# modified csv from ENSEMBL webpage
meta.1ktidy <- read.csv('00_util_scripts/ref/1kgp.2504sample.csv')

# convert 1kgp height snp vcf to gds ------
fetch.gds <- function(node, path) {node |> index.gdsn(path) |> read.gdsn()}

snpgdsVCF2GDS('~/append-ssd/work/vcf/height.ensb.1kgp.chr1.14.vcf.gz',
              "height.ensb.1kgp.chr1.14.gds")

gobj <- snpgdsOpen('height.ensb.1kgp.chr1.14.gds',
                   readonly = F)

sid.gds <- gobj |>
  fetch.gds('sample.id')

all(meta.1ktidy$sample == sid.gds)

## add sample annot (sex & pop) -----
to.gds <- meta.1ktidy |>
  dplyr::select(-sample)

gobj |>
  add.gdsn('sample.annot', to.gds)

# calc Fst -------
pop_code <- gobj |>
  fetch.gds("sample.annot/subpop")

pid.gds <- gobj |>
  fetch.gds('snp.id')

gobj.tb <- c('snp.id','snp.chromosome','snp.position') |>
  set_names() |>
  map(\(x)fetch.gds(gobj, path = x)) |>
  as_tibble()

# Two populations: CHS & CHB
flag <- pop_code %in% c("CHS", "CHB",'CDX')
samp.sel <- sid.gds[flag]
pop.sel <- pop_code[flag]
res.fst <- gobj |>
  snpgdsFst(sample.id=samp.sel, population=as.factor(pop.sel),
            method = 'W&C84', with.id = T)

res.fst$Fst        # Weir and Cockerham weighted Fst estimate
res.fst$MeanFst    # Weir and Cockerham mean Fst estimate
summary(res.fst$FstSNP)

tb.fst <- tibble(snp.id = res.fst$snp.id, fst = res.fst$FstSNP) |>
  left_join(gobj.tb)

tb.fst |>
  ggplot(aes(snp.position, fst)) +
  facet_wrap(~snp.chromosome, scales = 'free_x') +
  geom_point() +
  theme_pubr(x.text.angle = 45) +
  labs(title = 'Fst value in 1KGP Chinese populations')

tb.fst |>
  slice_max(fst, n = 6)

maf.candid2 <- gobj |>
  fetch.gds('genotype') |>
  as_tibble() |>
  dplyr::select(V970) |>
  bind_cols(meta.1ktidy) |>
  summarise(maf = 1 - mean(V970)/2, .by = subpop)

maf.candid2 |>
  mutate(Population.code = subpop,
         maf = signif(maf, 3)) |>
  left_join(geography_1kgp) |>
  plot_1kgp_globe('rs10800409')

gobj |> snpgdsClose()

# all chromosomes --------
seqVCF2GDS('~/append-ssd/work/vcf/height.ensb.1kgp.all2503.vcf.gz',
           "height.ensembl.gds") |>
  system.time()

gobj <- seqOpen('height.ensembl.gds')

gobj

sid2018 <- gobj |>
  seqGetData('sample.id')

meta2018 <- tibble(sample = sid2018) |>
  left_join(meta.1ktidy) |>
  mutate(subpop = ifelse(is.na(subpop), 'NA', subpop))

chn.id <- meta2018 |>
  filter(subpop %in% c('CHS','CDX','CHB')) |>
  pull(sample)

chn.pop <- meta2018 |>
  filter(subpop %in% c('CHS','CDX','CHB')) |>
  pull(subpop) |>
  as_factor()

## Fst by sliding window ----------
sld.wnd <- gobj |>
  snpgdsSlidingWindow(winsize=5e5, shift=1e5,
                      FUN="snpgdsFst", as.is="numeric",
                      population= chn.pop, sample.id = chn.id)

sld.wnd |>
  glimpse()

extract.chr <- function(x){
  chr.id <- sld.wnd[x] |> names() |>
  str_extract('\\d+')

sld.wnd[x:(x+2)] |>
  as_tibble() |>
  set_names(c('Fst','snp.count','position')) |>
  mutate(chr = chr.id)}

all.chr.fst <- seq.int(3,87,by = 4) |>
  map(extract.chr) |>
  list_rbind() |>
  filter(!is.na(Fst)) |>
  mutate(chr = as.integer(chr))

all.chr.fst |>
  ggplot(aes(position, Fst)) +
  geom_point(size = 1) +
  facet_grid(~chr, scales = 'free_x') +
  theme_pubr() +
  theme(axis.text.x = element_blank(),panel.spacing = unit(0,units = 'mm')) +
  labs(title = 'Published SNPs associated with body height in 1KGP Chinese populations')

all.chr.fst |>
  slice_max(Fst, n = 3)

## by individual SNP ----------
fst.bysnp <- gobj |>
  snpgdsFst(population = chn.pop, sample.id = chn.id,
            method = 'W&H02', with.id = T)

gobj.tb <- gobj |>
  seqGetData(c('variant.id','chromosome','position')) |>
  as_tibble()

fst.snp.tb <- fst.bysnp[c('snp.id','FstSNP')] |>
  as_tibble() |>
  rename(variant.id = snp.id) |>
  left_join(gobj.tb) |>
  mutate(chr = as.integer(chromosome))

fst.snp.tb |>
  ggplot(aes(position, FstSNP)) +
  geom_point(size = 1) +
  facet_grid(~chr, scales = 'free_x') +
  theme_pubr() +
  theme(axis.text.x = element_blank(),panel.spacing = unit(0,units = 'mm')) +
  labs(title = 'SNPs associated with body height in 1KGP Chinese populations')

fst.snp.tb |>
  slice_max(FstSNP, n = 3)

fst.snp.tb |> 
  select(-chromosome) |>
  write_csv('mission/exercise/height.ensembl.snp.chn.fst.csv')

# plot 2nd SNP MAF on globe -------
geno2b <- gobj |>
  seqGetData('$dosage')

geno2b |>
  glimpse()

df <- geno2b[,18279] |>
  as_tibble() |>
  bind_cols(meta2018) |>
  summarise(maf = mean(value)/2, .by = subpop)

df |>
  plot_1kgp_globe(pop_column = 'subpop', 'rs174547')

seqClose(gobj)
